import cv2
import numpy as np
from keras.models import load_model
import os
import tkinter as tk

dnnnet = cv2.dnn.readNetFromCaffe("deploy.prototxt", "res10_300x300_ssd_iter_140000.caffemodel")
emotion_model = load_model('model.h5')
class_labels = ['生气', '厌恶', '害怕', '开心', '悲伤', '惊讶', '正常']
# or
# 初始化截图计数器
screenshot_counter = 0
# 摄像头状态
camera_on = False
# dnnnet = cv2.dnn.readNetFromTensorflow("path/to/opencv_face_detector_uint8.pb", "path/to/opencv_face_detector.pbtxt")
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
cap.set(cv2.CAP_PROP_FPS, 60)
save_path = 'screen'

def close_camera():
    global cap
    if cap.isOpened():
        cap.release()


def open_screen_folder():
    screen_folder = os.path.join(os.getcwd(), "screen")
    if not os.path.exists(screen_folder):
        os.makedirs(screen_folder)
    os.startfile(screen_folder)


def toggle_camera():
    global camera_on
    if not camera_on:
        camera_on = True
        btn_open_camera.config(state="disabled")
        btn_close_camera.config(state="normal")
        open_camera()
    else:
        camera_on = False
        btn_open_camera.config(state="normal")
        btn_close_camera.config(state="disabled")
        close_camera()


def open_camera():
    global cap
    global screenshot_counter
    if not cap.isOpened():
        cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
        cap.set(cv2.CAP_PROP_FPS, 60)

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        h, w = frame.shape[:2]  # Get the frame dimensions
        blob = cv2.dnn.blobFromImage(frame, 1.0, (300, 300), (104., 177., 123.))  # Create blob from frame
        dnnnet.setInput(blob)
        detections = dnnnet.forward()  # Perform face detection

        for i in range(detections.shape[2]):
            confidence = detections[0, 0, i, 2]
            if confidence > 0.5:  # Adjust the confidence threshold as needed
                box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])  # Scale the box coordinates
                x_min, y_min, x_max, y_max = box.astype(int)
                face_roi = frame[y_min:y_max, x_min:x_max]
                face_roi = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
                face_roi = cv2.resize(face_roi, (48, 48))
                face_roi = face_roi.reshape(-1, 48, 48, 1)
                face_roi = face_roi / 255.0
                emotion_label = emotion_model.predict(face_roi)
                emotion_label = np.argmax(emotion_label)

                # Draw bounding box and label
                cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)

                if emotion_label == 0:
                    emotion_text = 'Angry'
                elif emotion_label == 1:
                    emotion_text = 'Disgust'
                elif emotion_label == 2:
                    emotion_text = 'Fear'
                elif emotion_label == 3:
                    emotion_text = 'Happy'
                elif emotion_label == 4:
                    emotion_text = 'Sad'
                elif emotion_label == 5:
                    emotion_text = 'Surprise'
                else:
                    emotion_text = 'Neutral'
                cv2.putText(frame, emotion_text, (x_min + 5, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

        cv2.imshow('DNN Face Detection', frame)

        # Press 'q' key to exit
        key = cv2.waitKey(1)
        if key & 0xFF == ord('s'):
            screenshot_filename = os.path.join(save_path, f"screenshot_{screenshot_counter}.png")
            cv2.imwrite(screenshot_filename, frame)
            print(f"Screenshot saved as {screenshot_filename}")
            screenshot_counter += 1
        # 检测按键，如果按下 'q' 键则退出循环
        elif key & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


# 创建主界面
root = tk.Tk()
root.title("Emotion Detection")
root.geometry("800x600")

# 创建按钮：打开摄像头
btn_open_camera = tk.Button(root, text="打开摄像头", command=toggle_camera)
btn_open_camera.pack(side=tk.LEFT, padx=30)

# 创建按钮：关闭摄像头
btn_close_camera = tk.Button(root, text="关闭摄像头", command=toggle_camera, state="disabled")
btn_close_camera.pack(side=tk.RIGHT, padx=30)

# 创建按钮：打开截图文件夹
btn_open_screen_folder = tk.Button(root, text="打开截图文件夹", command=open_screen_folder)
btn_open_screen_folder.pack(side=tk.RIGHT, padx=10)
# 运行主界面
root.mainloop()
